package test;

import com.gitee.mp.annotation.Collection;
import com.gitee.mp.annotation.Field;
import com.gitee.mp.domain.condition.ConditionBuilder;
import com.gitee.mp.domain.condition.PageInfo;
import com.gitee.mp.extend.SimilaritySetter;
import io.milvus.grpc.ConsistencyLevel;
import io.milvus.grpc.DataType;
import lombok.Data;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author hongda.li
 */
@Data
@Collection(name = "test_name",
        description = "This is a simple collection example",
        level = ConsistencyLevel.Session)
public class SimpleCollection implements SimilaritySetter {

    @Field(name = "id", type = DataType.VarChar, isPrimary = true)
    private String testId;

    @Field(name = "code", type = DataType.String)
    private String code;

    @Field(name = "test_content", type = DataType.String)
    private String content;

    @Field(name = "time", type = DataType.String)
    private String createdTime;

    @Field(name = "embedding", type = DataType.FloatVector)
    private List<Float> embedding;

    private Float distance;

    @Override
    public void setSimilarity(float similarity) {
        this.distance = similarity;
    }

    public static void main(String[] args) {
        ConditionBuilder.newBuilder(SimpleCollection.class)
                // 查询字段
                .select(SimpleCollection::getCode)
                // 查询字段
                .select(SimpleCollection::getTestId)
                // 检索条件 and distance > 2
                .andGt(SimpleCollection::getDistance, 2F)
                // 检索条件 and content == \"test\"
                .andEq(SimpleCollection::getContent, "test")
                // 检索条件 and code in [\"a\", \"b\", \"c\"]
                .andIn(SimpleCollection::getCode, Arrays.asList("a", "b", "c"))
                // 检索条件 or test_id != \"\"
                .orIsNotEmpty(SimpleCollection::getTestId)
                // 检索条件 or distance <= -2
                .orLte(SimpleCollection::getDistance, -2F)
                // 检索条件 limit 20 offset 10 * 20
                .pageInfo(new PageInfo().setPage(10L).setSize(20L))
                // 检索条件 offset 100
                .offset(100)
                // 单向量检索 search vector [[0.5]]
                .search(Collections.singletonList(0.5F));
    }
}
